/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * 本文件在MIT许可证条款下发布。
 *
 * 功能：
 * 本文件实现了LoongArch64架构下EFI启动时的内存管理单元(MMU)相关功能
 * 主要包含以下功能:
 * 1. 页表结构定义和页表项创建
 * 2. 物理内存到虚拟内存的映射
 * 3. EFI退出后的页表生成
 * 4. 内存映射的处理和内核参数设置
 * 5. MMU初始化和配置
 * 6. 内核参数的设置
 * 7. 内核参数的初始化
 * 8. 内核参数的验证
 * 9. 内核参数的打印
 * 10. 内核参数的保存
 * 11. 内核参数的加载
 * 12. 内核参数的转换
 *
 *
 * 作者：Amor，2025年4月
 */

#include <algorithm>

#include <kernel.h>
#include <arch_kernel.h>
#include <boot/platform.h>
#include <boot/stage2.h>
#include <efi/types.h>
#include <efi/boot-services.h>
#include <string.h>

#include "efi_platform.h"
#include "generic_mmu.h"
#include "mmu.h"

//#define TRACE_MMU
#ifdef TRACE_MMU
#	define TRACE(x...) dprintf(x)
#else
#	define TRACE(x...) ;
#endif

// Ignore memory above 512GB
#define PHYSICAL_MEMORY_LOW		0x00000000
#define PHYSICAL_MEMORY_HIGH	0x8000000000ull

// LoongArch64 MMU相关常量
#define PAGE_SIZE			4096
#define PAGE_SHIFT			12
#define PAGE_MASK			(PAGE_SIZE - 1)

// LoongArch64页表项定义
struct page_table_entry {
	uint64 valid : 1;      // 有效位
	uint64 dirty : 1;      // 脏位
	uint64 plv : 2;        // 特权级
	uint64 mat : 3;        // 内存属性
	uint64 global : 1;     // 全局位
	uint64 reserved : 2;   // 保留位
	uint64 ppn : 40;       // 物理页号
	uint64 reserved2 : 14; // 保留位
};

// 全局页表基址
static uint64* sPageDirectory = NULL;

// 从物理地址获取虚拟地址
static inline void*
VirtFromPhys(uint64_t physAdr)
{
	return (void*)physAdr;
}

// 创建页表项
static inline uint64
MakePageTableEntry(uint64 physAddr, uint64 flags)
{
	page_table_entry entry;
	memset(&entry, 0, sizeof(entry));

	entry.valid = 1;
	entry.dirty = (flags & TLBELO_D) ? 1 : 0;
	entry.plv = (flags & TLBELO_PLV) >> 2;
	entry.mat = (flags & TLBELO_MAT) >> 4;
	entry.global = (flags & TLBELO_G) ? 1 : 0;
	entry.ppn = (physAddr >> PAGE_SHIFT);

	return *(uint64*)&entry;
}

// 映射物理内存区域到虚拟地址空间
static void
MapRange(addr_t virtAdr, phys_addr_t physAdr, size_t size, uint64 flags)
{
	TRACE("MapRange(%#" B_PRIxADDR " - %#" B_PRIxADDR ", %#" B_PRIxADDR " - %#" B_PRIxADDR ", %#"
		B_PRIxADDR ")\n", virtAdr, virtAdr + size - 1, physAdr, physAdr + size - 1, size);

	for (addr_t offset = 0; offset < size; offset += PAGE_SIZE) {
		addr_t virtPage = virtAdr + offset;
		phys_addr_t physPage = physAdr + offset;

		// 计算页表索引
		uint32 pdpIndex = (virtPage >> 39) & 0x1ff;
		uint32 pdIndex = (virtPage >> 30) & 0x1ff;
		uint32 ptIndex = (virtPage >> 21) & 0x1ff;
		uint32 pIndex = (virtPage >> 12) & 0x1ff;

		// 确保页目录存在
		if (sPageDirectory[pdpIndex] == 0) {
			uint64* pd = (uint64*)mmu_allocate_page();
			if (pd == NULL)
				panic("无法分配页目录");
			memset(pd, 0, PAGE_SIZE);
			sPageDirectory[pdpIndex] = MakePageTableEntry((addr_t)pd, TLBELO_V);
		}

		// 获取页目录
		uint64* pd = (uint64*)VirtFromPhys(sPageDirectory[pdpIndex] & ~PAGE_MASK);

		// 确保页表存在
		if (pd[pdIndex] == 0) {
			uint64* pt = (uint64*)mmu_allocate_page();
			if (pt == NULL)
				panic("无法分配页表");
			memset(pt, 0, PAGE_SIZE);
			pd[pdIndex] = MakePageTableEntry((addr_t)pt, TLBELO_V);
		}

		// 获取页表
		uint64* pt = (uint64*)VirtFromPhys(pd[pdIndex] & ~PAGE_MASK);

		// 确保页表项存在
		if (pt[ptIndex] == 0) {
			uint64* p = (uint64*)mmu_allocate_page();
			if (p == NULL)
				panic("无法分配页");
			memset(p, 0, PAGE_SIZE);
			pt[ptIndex] = MakePageTableEntry((addr_t)p, TLBELO_V);
		}

		// 获取页
		uint64* p = (uint64*)VirtFromPhys(pt[ptIndex] & ~PAGE_MASK);

		// 设置页表项
		p[pIndex] = MakePageTableEntry(physPage, flags | TLBELO_V);
	}
}

// 生成EFI退出后使用的页表
uint64_t
arch_mmu_generate_post_efi_page_tables(size_t memory_map_size,
	efi_memory_descriptor* memory_map, size_t descriptor_size,
	uint32_t descriptor_version)
{
	// 分配页目录
	sPageDirectory = (uint64*)mmu_allocate_page();
	if (sPageDirectory == NULL)
		panic("无法分配页目录");
	memset(sPageDirectory, 0, PAGE_SIZE);

	// 遍历内存映射
	addr_t addr = (addr_t)memory_map;
	for (size_t i = 0; i < memory_map_size / descriptor_size; ++i) {
		efi_memory_descriptor* entry =
			(efi_memory_descriptor*)(addr + i * descriptor_size);

		// 跳过高于512GB的内存
		if (entry->PhysicalStart >= PHYSICAL_MEMORY_HIGH)
			continue;

		// 计算区域大小
		size_t size = entry->NumberOfPages * PAGE_SIZE;

		// 映射物理内存到相同的虚拟地址（恒等映射）
		MapRange(entry->PhysicalStart, entry->PhysicalStart, size,
			TLBELO_D | TLBELO_G | TLBELO_MAT);

		// 如果是加载器代码或数据，也映射到内核空间
		if (entry->Type == EfiLoaderCode || entry->Type == EfiLoaderData) {
			MapRange(KERNEL_BASE + entry->PhysicalStart, entry->PhysicalStart, size,
				TLBELO_D | TLBELO_G | TLBELO_MAT);
		}
	}

	// 返回页表基址
	return (uint64_t)sPageDirectory;
}

// EFI退出后的MMU设置
void
arch_mmu_post_efi_setup(size_t memory_map_size,
	efi_memory_descriptor* memory_map, size_t descriptor_size,
	uint32_t descriptor_version)
{
	// 更新内核参数中的物理内存映射
	gKernelArgs.physical_memory_range[0].start = PHYSICAL_MEMORY_LOW;
	gKernelArgs.physical_memory_range[0].size = PHYSICAL_MEMORY_HIGH - PHYSICAL_MEMORY_LOW;
	gKernelArgs.num_physical_memory_ranges = 1;

	// 设置内核参数中的虚拟分配范围
	gKernelArgs.virtual_allocated_range[0].start = KERNEL_BASE;
	gKernelArgs.virtual_allocated_range[0].size = KERNEL_SIZE;
	gKernelArgs.num_virtual_allocated_ranges = 1;

	// 处理内存映射
	addr_t addr = (addr_t)memory_map;
	for (size_t i = 0; i < memory_map_size / descriptor_size; ++i) {
		efi_memory_descriptor* entry =
			(efi_memory_descriptor*)(addr + i * descriptor_size);

		// 跳过高于512GB的内存
		if (entry->PhysicalStart >= PHYSICAL_MEMORY_HIGH)
			continue;

		// 计算区域大小
		size_t size = entry->NumberOfPages * PAGE_SIZE;

		switch (entry->Type) {
			case EfiLoaderCode:
			case EfiLoaderData:
			case EfiBootServicesCode:
			case EfiBootServicesData:
			case EfiConventionalMemory: {
				// 添加可用内存区域
				insert_physical_memory_range(entry->PhysicalStart, size);
				break;
			}

			case EfiACPIReclaimMemory: {
				// 添加ACPI可回收内存区域
				insert_physical_memory_range(entry->PhysicalStart, size);
				insert_physical_allocated_range(entry->PhysicalStart, size);
				break;
			}

			case EfiRuntimeServicesCode:
			case EfiRuntimeServicesData: {
				// 添加EFI运行时服务区域
				insert_physical_allocated_range(entry->PhysicalStart, size);
				break;
			}

			case EfiReservedMemoryType:
			case EfiUnusableMemory:
			case EfiMemoryMappedIO:
			case EfiMemoryMappedIOPortSpace:
			case EfiPalCode: {
				// 添加保留内存区域
				insert_physical_allocated_range(entry->PhysicalStart, size);
				break;
			}

			case EfiACPIMemoryNVS: {
				// 添加ACPI NVS内存区域
				insert_physical_allocated_range(entry->PhysicalStart, size);
				break;
			}
		}
	}

	// 设置内核参数中的物理内存映射区域
	gKernelArgs.arch_args.physMap.start = KERNEL_BASE;
	gKernelArgs.arch_args.physMap.size = PHYSICAL_MEMORY_HIGH - PHYSICAL_MEMORY_LOW;

	// 设置内核参数中的虚拟保留区域
	gKernelArgs.arch_args.virtual_ranges_to_keep[0].start = KERNEL_BASE;
	gKernelArgs.arch_args.virtual_ranges_to_keep[0].size = KERNEL_SIZE;
	gKernelArgs.arch_args.num_virtual_ranges_to_keep = 1;
}